package com.cqs.leetcode;

import java.util.*;

class RandomizedCollection {


    /**
     * 存储元素
     */
    private List<Integer> elems = new ArrayList<>();


    private Map<Integer, Set<Integer>> indexs = new HashMap<>();

    /**
     * Initialize your data structure here.
     */
    public RandomizedCollection() {
    }

    /**
     * Inserts a value to the collection. Returns true if the collection did not already contain the specified element.
     */
    public boolean insert(int val) {
        elems.add(val);
        //添加到indexs中
        Set<Integer> list = indexs.getOrDefault(val, new HashSet<>());
        list.add(elems.size() - 1);
        indexs.put(val, list);
        return list.size() == 1;
    }

    /**
     * Removes a value from the collection. Returns true if the collection contained the specified element.
     */
    public boolean remove(int val) {
        if (!indexs.containsKey(val)) return false;
        Set<Integer> del = indexs.get(val);
        Iterator<Integer> iterator = del.iterator();
        Integer delIdx = iterator.next();
        iterator.remove();

        //数组最后一个元素移动到delIdx位置
        Integer lastEle = elems.get(elems.size() - 1);
        elems.set(delIdx, lastEle);
        elems.remove(elems.size() - 1);
        //更新最后要给元素在Map的映射关系
        Set<Integer> set = indexs.get(lastEle);
        set.add(delIdx);
        set.remove(elems.size());
        if (del.size() == 0) indexs.remove(val);
        return true;
    }

    /**
     * Get a random element from the collection.
     */
    public int getRandom() {
        int idx = new Random().nextInt(elems.size());
        return elems.size() == 0 ? 0 : elems.get(idx);
    }

    public static void main(String[] args) {
        RandomizedCollection collection = new RandomizedCollection();
        collection.insert(1);
        collection.insert(1);
        collection.insert(2);
        collection.insert(2);
        collection.insert(2);
        collection.remove(1);
        collection.remove(1);
        collection.remove(2);
        collection.insert(1);
        collection.remove(2);
        System.out.println(collection.indexs);
    }

}
